import torch
# 初始化模型参数
batch_size = 5    # 批量大小
seq_len = 3       # 序列长度
input_size = 4    # 特征维度
hidden_size = 2   # 隐藏层大小
cell = torch.nn.RNNCell(input_size=input_size,hidden_size=hidden_size)
# 输入张量是三维的,维度为(seq, batch, features)
dataset = torch.randn(seq_len, batch_size, input_size)
hidden0 = torch.zeros(batch_size, hidden_size) # 因为是单隐藏,所以没有layer_size
# 由于没有先验,所以初始化第一个隐藏层单元设置为全零
print("=" * 30, 0, "=" * 30)
input0 = dataset[0,:,:]
print("Input0 Size:", input0.shape)
# 隐藏状态
hidden1 = cell(input0, hidden0) # hidden1:(5,2)
print("outputs1 size:", hidden1.shape)
print(hidden1)

print("=" * 30, 1, "=" * 30)
input1= dataset[1,:,:]
print("Input1 Size:", input1.shape)
# 隐藏状态
hidden2 = cell(input1, hidden1) # hidden2:(5,2)
print("outputs2 size:", hidden2.shape)
print(hidden2)

print("=" * 30, 2, "=" * 30)
input2= dataset[2,:,:]
print("Input2 Size:", input2.shape)
# 隐藏状态
hidden3 = cell(input2, hidden2) # hidden3:(5,2)
print("outputs3 size:", hidden3.shape)
print(hidden3)

#-------------或者用循环实现-------------------
for idx, input in enumerate(dataset):
    # 每次取出一个seq:(batch, features)
    print("=" * 30, idx, "=" * 30)
    print("Input Size:", input.shape)    
    # 迭代隐藏状态
    hidden0 = cell(input, hidden0)
    print("outputs size:", hidden0.shape)
    print(hidden0) 
